From 53297aaac05d2ea79cf9ca8bd1f9c3d0fbd51119 Mon Sep 17 00:00:00 2001
From: ShiroKSH <kushidashiro@gmail.com>
Date: Wed, 15 Jul 2026 13:58:30 +0300
Subject: [PATCH] fix: preserve component band offsets

---
 python/sedona/spark/raster/sample_model.py |  4 +-
 python/tests/raster/test_sample_model.py   | 47 ++++++++++++++++++++++
 2 files changed, 49 insertions(+), 2 deletions(-)
 create mode 100644 python/tests/raster/test_sample_model.py

diff --git a/python/sedona/spark/raster/sample_model.py b/python/sedona/spark/raster/sample_model.py
index fdf30f38397..5b5ec222247 100644
--- a/python/sedona/spark/raster/sample_model.py
+++ b/python/sedona/spark/raster/sample_model.py
@@ -80,9 +80,9 @@ def as_numpy(self, data_buffer: DataBuffer) -> np.ndarray:
         if self.scanline_stride == self.width and self.pixel_stride == 1:
             # Fast path: no gaps between pixels
             band_arrs = []
-            for bank_index in self.bank_indices:
+            for k, bank_index in enumerate(self.bank_indices):
                 bank_data = data_buffer.bank_data[bank_index]
-                offset = self.band_offsets[bank_index]
+                offset = self.band_offsets[k]
                 if offset != 0:
                     bank_data = bank_data[offset : (offset + self.width * self.height)]
                 band_arr = bank_data.reshape(self.height, self.width)
diff --git a/python/tests/raster/test_sample_model.py b/python/tests/raster/test_sample_model.py
new file mode 100644
index 00000000000..1403c12fac1
--- /dev/null
+++ b/python/tests/raster/test_sample_model.py
@@ -0,0 +1,47 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import numpy as np
+
+from sedona.spark.raster.data_buffer import DataBuffer
+from sedona.spark.raster.sample_model import ComponentSampleModel
+
+
+def test_component_sample_model_fast_path_uses_band_offsets() -> None:
+    sample_model = ComponentSampleModel(
+        DataBuffer.TYPE_INT,
+        2,
+        2,
+        1,
+        2,
+        [1, 0],
+        [3, 1],
+    )
+    data_buffer = DataBuffer(
+        DataBuffer.TYPE_INT,
+        [
+            np.array([10, 11, 12, 13, 14, 15, 16]),
+            np.array([20, 21, 22, 23, 24, 25, 26]),
+        ],
+        7,
+        [0, 0],
+    )
+
+    np.testing.assert_array_equal(
+        sample_model.as_numpy(data_buffer),
+        np.array([[[23, 24], [25, 26]], [[11, 12], [13, 14]]]),
+    )