diff --git a/projects/ng-table-virtual-scroll/src/lib/table-item-size.directive.ts b/projects/ng-table-virtual-scroll/src/lib/table-item-size.directive.ts
index a0cc79a..0e97161 100644
--- a/projects/ng-table-virtual-scroll/src/lib/table-item-size.directive.ts
+++ b/projects/ng-table-virtual-scroll/src/lib/table-item-size.directive.ts
@@ -1,5 +1,9 @@
 import { VIRTUAL_SCROLL_STRATEGY } from '@angular/cdk/scrolling';
-import { CanStick, CdkTable } from '@angular/cdk/table';
+import {
+  CanStick,
+  CdkTable,
+  CdkTableDataSourceInput,
+} from '@angular/cdk/table';
 import {
   AfterContentInit,
   ContentChild,
@@ -8,13 +12,23 @@ import {
   Input,
   NgZone,
   OnChanges,
-  OnDestroy
+  OnDestroy,
 } from '@angular/core';
 import { MatTable } from '@angular/material/table';
-import { combineLatest, from, Subject } from 'rxjs';
-import { delayWhen, distinctUntilChanged, map, startWith, switchMap, take, takeUntil, tap } from 'rxjs/operators';
 import { FixedSizeTableVirtualScrollStrategy } from './fixed-size-table-virtual-scroll-strategy';
 import { CdkTableVirtualScrollDataSource, isTVSDataSource, TableVirtualScrollDataSource } from './table-data-source';
+import { combineLatest, from, Subject } from 'rxjs';
+import {
+  delayWhen,
+  distinctUntilChanged,
+  map,
+  startWith,
+  switchMap,
+  take,
+  takeUntil,
+  tap,
+} from 'rxjs/operators';
+
 
 export function _tableVirtualScrollDirectiveStrategyFactory(tableDir: TableItemSizeDirective) {
   return tableDir.scrollStrategy;
@@ -61,7 +75,9 @@ const defaults = {
     deps: [forwardRef(() => TableItemSizeDirective)]
   }]
 })
-export class TableItemSizeDirective<T = unknown> implements OnChanges, AfterContentInit, OnDestroy {
+export class TableItemSizeDirective<T = unknown>
+  implements OnChanges, AfterContentInit, OnDestroy
+{
   private destroyed$ = new Subject<void>();
 
   // eslint-disable-next-line @angular-eslint/no-input-rename
@@ -84,21 +100,20 @@ export class TableItemSizeDirective<T = unknown> implements OnChanges, AfterCont
   bufferMultiplier: string | number = defaults.bufferMultiplier;
 
   @ContentChild(CdkTable, { static: false })
-  table: CdkTable<T>;
+  table!: CdkTable<T>;
 
   scrollStrategy = new FixedSizeTableVirtualScrollStrategy();
 
   dataSourceChanges = new Subject<void>();
 
-  private stickyPositions: Map<HTMLElement, number>;
+  private stickyPositions: Map<HTMLElement, number> | null = null;
   private resetStickyPositions = new Subject<void>();
   private stickyEnabled = {
     header: false,
-    footer: false
+    footer: false,
   };
 
-  constructor(private zone: NgZone) {
-  }
+  constructor(private zone: NgZone) {}
 
   ngOnDestroy() {
     this.destroyed$.next();
@@ -108,14 +123,19 @@ export class TableItemSizeDirective<T = unknown> implements OnChanges, AfterCont
 
   ngAfterContentInit() {
     const switchDataSourceOrigin = this.table['_switchDataSource'];
-    this.table['_switchDataSource'] = (dataSource: any) => {
+    this.table['_switchDataSource'] = (
+      dataSource:
+        | TableVirtualScrollDataSource<T>
+        | CdkTableVirtualScrollDataSource<T>,
+    ) => {
       switchDataSourceOrigin.call(this.table, dataSource);
       this.connectDataSource(dataSource);
     };
 
     const updateStickyColumnStylesOrigin = this.table.updateStickyColumnStyles;
     this.table.updateStickyColumnStyles = () => {
-      const stickyColumnStylesNeedReset = this.table['_stickyColumnStylesNeedReset'];
+      const stickyColumnStylesNeedReset =
+        this.table['_stickyColumnStylesNeedReset'];
       updateStickyColumnStylesOrigin.call(this.table);
       if (stickyColumnStylesNeedReset) {
         this.resetStickyPositions.next();
@@ -131,12 +151,10 @@ export class TableItemSizeDirective<T = unknown> implements OnChanges, AfterCont
         delayWhen(() => this.getScheduleObservable()),
         tap(() => {
           this.stickyPositions = null;
-        })
-      )
+        }),
+      ),
     ])
-      .pipe(
-        takeUntil(this.destroyed$)
-      )
+      .pipe(takeUntil(this.destroyed$))
       .subscribe(([stickyOffset]) => {
         if (!this.stickyPositions) {
           this.initStickyPositions();
@@ -150,37 +168,52 @@ export class TableItemSizeDirective<T = unknown> implements OnChanges, AfterCont
       });
   }
 
-  connectDataSource(dataSource: unknown) {
+  connectDataSource(
+    dataSource:
+      | TableVirtualScrollDataSource<T>
+      | CdkTableVirtualScrollDataSource<T>
+      | CdkTableDataSourceInput<T>,
+  ) {
     this.dataSourceChanges.next();
     if (!isTVSDataSource(dataSource)) {
-      throw new Error('[tvsItemSize] requires TableVirtualScrollDataSource or CdkTableVirtualScrollDataSource be set as [dataSource] of the table');
+      throw new Error(
+        '[tvsItemSize] requires TableVirtualScrollDataSource or CdkTableVirtualScrollDataSource be set as [dataSource] of the table',
+      );
     }
-    if (isMatTable(this.table) && !(dataSource instanceof TableVirtualScrollDataSource)) {
-      throw new Error('[tvsItemSize] requires TableVirtualScrollDataSource be set as [dataSource] of [mat-table]');
+    if (
+      isMatTable(this.table) &&
+      !(dataSource instanceof TableVirtualScrollDataSource)
+    ) {
+      throw new Error(
+        '[tvsItemSize] requires TableVirtualScrollDataSource be set as [dataSource] of [mat-table]',
+      );
     }
-    if (isCdkTable(this.table) && !(dataSource instanceof CdkTableVirtualScrollDataSource)) {
-      throw new Error('[tvsItemSize] requires CdkTableVirtualScrollDataSource be set as [dataSource] of [cdk-table]');
+    if (
+      isCdkTable(this.table) &&
+      !(dataSource instanceof CdkTableVirtualScrollDataSource)
+    ) {
+      throw new Error(
+        '[tvsItemSize] requires CdkTableVirtualScrollDataSource be set as [dataSource] of [cdk-table]',
+      );
     }
 
-    dataSource
-      .dataToRender$
+    dataSource.dataToRender$
       .pipe(
         distinctUntilChanged(),
         takeUntil(this.dataSourceChanges),
         takeUntil(this.destroyed$),
-        tap(data => this.scrollStrategy.dataLength = data.length),
-        switchMap(data =>
-          this.scrollStrategy
-            .renderedRangeStream
-            .pipe(
-              map(({
-                     start,
-                     end
-                   }) => typeof start !== 'number' || typeof end !== 'number' ? data : data.slice(start, end))
-            )
-        )
+        tap((data) => (this.scrollStrategy.dataLength = data.length)),
+        switchMap((data) =>
+          this.scrollStrategy.renderedRangeStream.pipe(
+            map(({ start, end }) =>
+              typeof start !== 'number' || typeof end !== 'number'
+                ? data
+                : data.slice(start, end),
+            ),
+          ),
+        ),
       )
-      .subscribe(data => {
+      .subscribe((data) => {
         this.zone.run(() => {
           dataSource.dataOfRange$.next(data);
         });
@@ -190,9 +223,13 @@ export class TableItemSizeDirective<T = unknown> implements OnChanges, AfterCont
   ngOnChanges() {
     const config = {
       rowHeight: +this.rowHeight || defaults.rowHeight,
-      headerHeight: this.headerEnabled ? +this.headerHeight || defaults.headerHeight : 0,
-      footerHeight: this.footerEnabled ? +this.footerHeight || defaults.footerHeight : 0,
-      bufferMultiplier: +this.bufferMultiplier || defaults.bufferMultiplier
+      headerHeight: this.headerEnabled
+        ? +this.headerHeight || defaults.headerHeight
+        : 0,
+      footerHeight: this.footerEnabled
+        ? +this.footerHeight || defaults.footerHeight
+        : 0,
+      bufferMultiplier: +this.bufferMultiplier || defaults.bufferMultiplier,
     };
     this.scrollStrategy.setConfig(config);
   }
@@ -201,43 +238,62 @@ export class TableItemSizeDirective<T = unknown> implements OnChanges, AfterCont
     if (!this.scrollStrategy.viewport) {
       this.stickyEnabled = {
         header: false,
-        footer: false
+        footer: false,
       };
-      return;
+      return false;
     }
 
-    const isEnabled = (rowDefs: CanStick[]) => rowDefs
-      .map(def => def.sticky)
-      .reduce((prevState, state) => prevState && state, true);
+    const isEnabled = (rowDefs: CanStick[]) =>
+      rowDefs
+        .map((def) => def.sticky)
+        .reduce((prevState, state) => prevState && state, true);
 
     this.stickyEnabled = {
       header: isEnabled(this.table['_headerRowDefs']),
       footer: isEnabled(this.table['_footerRowDefs']),
     };
+    return true;
   }
 
   private setStickyHeader(offset: number) {
-    this.scrollStrategy.viewport.elementRef.nativeElement.querySelectorAll(stickyHeaderSelector)
-      .forEach((el: HTMLElement) => {
+    let stickyOffset = offset;
+    this.scrollStrategy.viewport.elementRef.nativeElement
+      .querySelectorAll(stickyHeaderSelector)
+      .forEach((el: Element) => {
         const parent = el.parentElement;
+        if (!parent) return;
         let baseOffset = 0;
-        if (this.stickyPositions.has(parent)) {
-          baseOffset = this.stickyPositions.get(parent);
+        if (this.stickyPositions?.has(parent)) {
+          baseOffset = this.stickyPositions.get(parent)!;
         }
-        el.style.top = `${baseOffset - offset}px`;
+        el.setAttribute(
+          'style',
+          `${el.getAttribute('style')} top: ${baseOffset + offset + stickyOffset}px`,
+        );
+        stickyOffset += el.getBoundingClientRect().height;
       });
   }
 
   private setStickyFooter(offset: number) {
-    this.scrollStrategy.viewport.elementRef.nativeElement.querySelectorAll(stickyFooterSelector)
-      .forEach((el: HTMLElement) => {
-        const parent = el.parentElement;
-        let baseOffset = 0;
-        if (this.stickyPositions.has(parent)) {
-          baseOffset = this.stickyPositions.get(parent);
-        }
-        el.style.bottom = `${-baseOffset + offset}px`;
-      });
+    let stickyOffset = offset;
+    const elements = Array.from(
+      this.scrollStrategy.viewport.elementRef.nativeElement.querySelectorAll(
+        stickyFooterSelector,
+      ),
+    ).reverse();
+
+    elements.forEach((el: Element) => {
+      const parent = el.parentElement;
+      if (!parent) return;
+      let baseOffset = 0;
+      if (this.stickyPositions?.has(parent)) {
+        baseOffset = this.stickyPositions.get(parent)!;
+      }
+      el.setAttribute(
+        'style',
+        `${el.getAttribute('style')} bottom: ${-baseOffset + offset + stickyOffset}px`,
+      );
+    });
   }
 
   private initStickyPositions() {
@@ -246,27 +302,30 @@ export class TableItemSizeDirective<T = unknown> implements OnChanges, AfterCont
     this.setStickyEnabled();
 
     if (this.stickyEnabled.header) {
-      this.scrollStrategy.viewport.elementRef.nativeElement.querySelectorAll(stickyHeaderSelector)
-        .forEach(el => {
+      this.scrollStrategy.viewport.elementRef.nativeElement
+        .querySelectorAll(stickyHeaderSelector)
+        .forEach((el) => {
           const parent = el.parentElement;
-          if (!this.stickyPositions.has(parent)) {
-            this.stickyPositions.set(parent, parent.offsetTop);
+          if (!parent) return;
+          if (!this.stickyPositions?.has(parent)) {
+            this.stickyPositions?.set(parent, parent.offsetTop);
           }
         });
     }
 
     if (this.stickyEnabled.footer) {
-      this.scrollStrategy.viewport.elementRef.nativeElement.querySelectorAll(stickyFooterSelector)
-        .forEach(el => {
+      this.scrollStrategy.viewport.elementRef.nativeElement
+        .querySelectorAll(stickyFooterSelector)
+        .forEach((el) => {
           const parent = el.parentElement;
-          if (!this.stickyPositions.has(parent)) {
-            this.stickyPositions.set(parent, -parent.offsetTop);
+          if (!parent) return;
+          if (!this.stickyPositions?.has(parent)) {
+            this.stickyPositions?.set(parent, -parent.offsetTop);
           }
         });
     }
   }
 
-
   private getScheduleObservable() {
     // Use onStable when in the context of an ongoing change detection cycle so that we
     // do not accidentally trigger additional cycles.